{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "v_eWsiiMfu2q",
      "metadata": {
        "id": "v_eWsiiMfu2q"
      },
      "source": [
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/cohere-ai/notebooks/blob/main/notebooks/llmu/RAG_with_Connectors.ipynb\">\n",
        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f9e30089",
      "metadata": {
        "id": "f9e30089"
      },
      "source": [
        "# RAG with Connectors\n",
        "\n",
        "This notebook shows how to build a RAG-powered chatbot with Cohere's Chat endpoint using connectors. The chatbot can extract relevant information from external documents and produce verifiable, inline citations in its responses.\n",
        "\n",
        "Connectors are ways of connecting to data sources. These data sources could be internal documents, document databases, the broader internet, or any other source of context which can inform the replies generated by the model.\n",
        "\n",
        "We'll use the web search connector, a Cohere-managed connector that you can use without additional setup.\n",
        "\n",
        "The diagram below provides an overview of what we’ll build."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "Gg0u4eKVubYI",
      "metadata": {
        "id": "Gg0u4eKVubYI"
      },
      "source": [
        "![rag-workflow-3.png]()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5IdmW-I9NXpq",
      "metadata": {
        "id": "5IdmW-I9NXpq"
      },
      "source": [
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "r2jcKQ6iLefn",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "r2jcKQ6iLefn",
        "outputId": "acb51f35-43a2-4567-d8e6-913f00d57df6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m117.2/117.2 kB\u001b[0m \u001b[31m1.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m75.6/75.6 kB\u001b[0m \u001b[31m3.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m77.8/77.8 kB\u001b[0m \u001b[31m4.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.3/58.3 kB\u001b[0m \u001b[31m4.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "! pip install cohere -q"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "90f134ba",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        },
        "id": "90f134ba",
        "outputId": "f2236cef-f274-4100-dbcd-333b826f5ee8"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "import uuid\n",
        "import cohere\n",
        "from cohere import ChatConnector\n",
        "from typing import List\n",
        "\n",
        "co = cohere.Client(\"COHERE_API_KEY\") # Get your API key here: https://dashboard.cohere.com/api-keys"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "EavECqgqNJ8g",
      "metadata": {
        "cellView": "form",
        "id": "EavECqgqNJ8g"
      },
      "outputs": [],
      "source": [
        "#@title Enable text wrapping in Google Colab\n",
        "\n",
        "from IPython.display import HTML, display\n",
        "\n",
        "def set_css():\n",
        "  display(HTML('''\n",
        "  <style>\n",
        "    pre {\n",
        "        white-space: pre-wrap;\n",
        "    }\n",
        "  </style>\n",
        "  '''))\n",
        "get_ipython().events.register('pre_run_cell', set_css)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d319ece1",
      "metadata": {
        "id": "d319ece1"
      },
      "source": [
        "# Create a chatbot\n",
        "\n",
        "In connector mode, most of the implementation is taken care of by the endpoint, including deciding whether to retrieve information, generating queries, retrieving documents, chunking and reranking documents (post-retrieval), and generating the response. This greatly simplifies our code.\n",
        "\n",
        "The `Chatbot` class below handles the interaction between the user and chatbot.  We define the connector for the chatbot to use with the attribute `self.connectors`. In this notebook, we will use Cohere's `“web-search”` connector, which runs searches against a browser in safe mode.\n",
        "\n",
        "The run() method contains the logic for getting the user message, displaying the chatbot response with citations, along with a way for the user to end the conversation.\n",
        "\n",
        "Then, the chatbot responds to the user message.  We call `co.chat()` and supply a `connectors` parameter to make the chatbot component use connector mode.  All of the remaining implementation is taken care of by the endpoint, up to generating the response.\n",
        "\n",
        "We also pass the `conversation_id` parameter, which retains the interactions between the user and the chatbot in the same conversation thread. We enable the `stream` parameter so we can stream the chatbot response.\n",
        "\n",
        "We then print the chatbot's response.  In the case that the external information was used to generate a response, we also display documents and in-line citations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e52d521d",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        },
        "id": "e52d521d",
        "outputId": "b0f90f1c-17c8-46fa-d471-b11059767ede"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "class Chatbot:\n",
        "    def __init__(self, connectors: List[str]):\n",
        "        \"\"\"\n",
        "        Initializes an instance of the Chatbot class.\n",
        "\n",
        "        \"\"\"\n",
        "        self.conversation_id = str(uuid.uuid4())\n",
        "        self.connectors = [ChatConnector(id=connector) for connector in connectors]\n",
        "\n",
        "    def run(self):\n",
        "        \"\"\"\n",
        "        Runs the chatbot application.\n",
        "\n",
        "        \"\"\"\n",
        "        while True:\n",
        "            # Get the user message\n",
        "            message = input(\"User: \")\n",
        "\n",
        "            # Typing \"quit\" ends the conversation\n",
        "            if message.lower() == \"quit\":\n",
        "                print(\"Ending chat.\")\n",
        "                break\n",
        "            # else:                         # Uncomment for Google Colab to avoid printing the same thing twice\n",
        "            #     print(f\"User: {message}\") # Uncomment for Google Colab to avoid printing the same thing twice\n",
        "\n",
        "            # Generate response\n",
        "            response = co.chat_stream(\n",
        "                    message=message,\n",
        "                    model=\"command-r\",\n",
        "                    conversation_id=self.conversation_id,\n",
        "                    connectors=self.connectors,\n",
        "            )\n",
        "\n",
        "            # Print the chatbot response, citations, and documents\n",
        "            print(\"\\nChatbot:\")\n",
        "            citations = []\n",
        "            cited_documents = []\n",
        "\n",
        "            # Display response\n",
        "            for event in response:\n",
        "                if event.event_type == \"text-generation\":\n",
        "                    print(event.text, end=\"\")\n",
        "                elif event.event_type == \"citation-generation\":\n",
        "                    citations.extend(event.citations)\n",
        "                elif event.event_type == \"search-results\":\n",
        "                    cited_documents = event.documents\n",
        "\n",
        "            # Display citations and source documents\n",
        "            if citations:\n",
        "              print(\"\\n\\nCITATIONS:\")\n",
        "              for citation in citations:\n",
        "                print(citation)\n",
        "\n",
        "              print(\"\\nDOCUMENTS:\")\n",
        "              for document in cited_documents:\n",
        "                print({'id': document['id'],\n",
        "                      'snippet': document['snippet'][:50] + '...',\n",
        "                      'title': document['title'],\n",
        "                      'url': document['url']})\n",
        "\n",
        "            print(f\"\\n{'-'*100}\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "rJbwGPksPfh0",
      "metadata": {
        "id": "rJbwGPksPfh0"
      },
      "source": [
        "# Run the chatbot"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c755140c",
      "metadata": {
        "id": "c755140c"
      },
      "source": [
        "We can now run the chatbot.  For this, we create the instance of `Chatbot` using Cohere's managed web-search connector.  Then we run the chatbot by invoking the `run()` method.\n",
        "\n",
        "The format of each citation is:\n",
        "- `start`: The starting point of a span where one or more documents are referenced\n",
        "- `end`: The ending point of a span where one or more documents are referenced\n",
        "- `text`: The text representing this span\n",
        "- `document_ids`: The IDs of the documents being referenced (`doc_0` being the ID of the first document passed to the `documents` creating parameter in the endpoint call, and so on)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "99e5005b",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 999
        },
        "id": "99e5005b",
        "outputId": "4609e72c-df6f-4c77-8132-cc0e73b80eee"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "User: What is Cohere's LLM University\n",
            "\n",
            "Chatbot:\n",
            "LLM University, offered by Cohere, is a set of comprehensive learning resources for anyone interested in Natural Language Processing (NLP), from beginners to advanced learners. The curriculum aims to provide a solid foundation in NLP and equips learners with the skills needed to develop their own AI applications.\n",
            "The course covers various topics, including semantic search, generation, classification, embeddings, and other NLP techniques. Learners can explore these concepts through hands-on exercises and practical code examples.\n",
            "Join the Discord community to connect with other learners and access the latest updates!\n",
            "\n",
            "CITATIONS:\n",
            "start=27 end=33 text='Cohere' document_ids=['web-search_0', 'web-search_1']\n",
            "start=47 end=79 text='comprehensive learning resources' document_ids=['web-search_1']\n",
            "start=105 end=132 text='Natural Language Processing' document_ids=['web-search_0', 'web-search_1']\n",
            "start=133 end=138 text='(NLP)' document_ids=['web-search_0', 'web-search_1']\n",
            "start=145 end=176 text='beginners to advanced learners.' document_ids=['web-search_0', 'web-search_1']\n",
            "start=181 end=191 text='curriculum' document_ids=['web-search_0', 'web-search_1']\n",
            "start=210 end=233 text='solid foundation in NLP' document_ids=['web-search_0', 'web-search_1']\n",
            "start=263 end=314 text='skills needed to develop their own AI applications.' document_ids=['web-search_0', 'web-search_1']\n",
            "start=359 end=374 text='semantic search' document_ids=['web-search_0', 'web-search_1']\n",
            "start=376 end=386 text='generation' document_ids=['web-search_0', 'web-search_1']\n",
            "start=388 end=402 text='classification' document_ids=['web-search_0', 'web-search_1']\n",
            "start=404 end=414 text='embeddings' document_ids=['web-search_0', 'web-search_1']\n",
            "start=420 end=441 text='other NLP techniques.' document_ids=['web-search_0', 'web-search_1']\n",
            "start=486 end=504 text='hands-on exercises' document_ids=['web-search_0', 'web-search_1']\n",
            "start=509 end=533 text='practical code examples.' document_ids=['web-search_0', 'web-search_1']\n",
            "start=543 end=560 text='Discord community' document_ids=['web-search_0', 'web-search_1']\n",
            "\n",
            "DOCUMENTS:\n",
            "{'id': 'web-search_0', 'snippet': 'Guides and ConceptsAPI ReferenceRelease NotesAppli...', 'title': 'LLM University (LLMU) | Cohere', 'url': 'https://docs.cohere.com/docs/llmu'}\n",
            "{'id': 'web-search_1', 'snippet': 'Introducing LLM University — Your Go-To Learning R...', 'title': 'Introducing LLM University — Your Go-To Learning Resource for NLP🎓', 'url': 'https://txt.cohere.com/llm-university/'}\n",
            "{'id': 'web-search_2', 'snippet': 'Skip to main content\\n\\nMadras High Court Reads Down...', 'title': 'LawBeat | Madras High Court Reads Down University Admission Rule Mandating 2-Yr LLM for PhD Admission', 'url': 'https://lawbeat.in/news-updates/madras-high-court-reads-down-university-admission-rule-mandating-2-yr-llm-phd-admission'}\n",
            "{'id': 'web-search_3', 'snippet': 'Take your legal expertise to the next level with a...', 'title': 'LLM Program', 'url': 'https://www.law.umaryland.edu/academics/llm-program/'}\n",
            "{'id': 'web-search_4', 'snippet': \"The People's Network\\n\\nSign In with Facebook\\n\\nBy cl...\", 'title': 'Revolutionizing AI: University of Michigan and Apple Team Up to Boost LLM Efficiency', 'url': 'https://bnnbreaking.com/world/us/revolutionizing-ai-university-of-michigan-and-apple-team-up-to-boost-llm-efficiency'}\n",
            "{'id': 'web-search_5', 'snippet': 'Ministers urged to tackle “damaging” trial delays ...', 'title': 'LLM Master of Laws (General) Degree | University of Law', 'url': 'https://www.law.ac.uk/study/postgraduate/law/llm-master-of-laws-general/'}\n",
            "{'id': 'web-search_6', 'snippet': \"The People's Network\\n\\nSign In with Facebook\\n\\nBy cl...\", 'title': \"Tsinghua University's Ouroboros Framework: Revolutionizing LLM Inference Speed by 2.8x\", 'url': 'https://bnnbreaking.com/world/china/tsinghua-universitys-ouroboros-framework-revolutionizing-llm-inference-speed-by-28x'}\n",
            "{'id': 'web-search_7', 'snippet': 'Skip to main content\\n\\nSupport the Law School\\n\\nCons...', 'title': 'LLM & Graduate Programs • Graduate Admissions • Penn Carey Law', 'url': 'https://www.law.upenn.edu/admissions/grad/'}\n",
            "{'id': 'web-search_8', 'snippet': 'Skip to navigation | Skip to main content | Skip t...', 'title': 'LLM Law (2024 entry) | The University of Manchester', 'url': 'https://www.manchester.ac.uk/study/masters/courses/list/08446/llm-law/'}\n",
            "\n",
            "----------------------------------------------------------------------------------------------------\n",
            "\n",
            "User: quit\n",
            "Ending chat.\n"
          ]
        }
      ],
      "source": [
        "# Define the connector\n",
        "connectors = [\"web-search\"]\n",
        "\n",
        "# Create an instance of the Chatbot class\n",
        "chatbot = Chatbot(connectors)\n",
        "\n",
        "# Run the chatbot\n",
        "chatbot.run()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7BOGEfHGPrCX",
      "metadata": {
        "id": "7BOGEfHGPrCX"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.5"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
